import os
import re
import docx
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from datetime import datetime
import time
from collections import defaultdict

# 升级方舟 SDK 到最新版本：pip install -U 'volcengine-python-sdk[ark]'
from volcenginesdkarkruntime import Ark

# 方舟API配置
ARK_MODEL = "deepseek-v3-1-250821"  # 模型名称
# 从环境变量获取API Key，也可直接替换为字符串（不推荐）
#ARK_API_KEY = os.environ.get("ARK_API_KEY")
ARK_API_KEY ="2c80c785-0460-4707-a917-10b460608862"
##print("API_KEY:"+os.environ.get("ARK_API_KEY"))

class ArkAPIClient:
    """方舟SDK客户端工具类"""
    
    def __init__(self, api_key=ARK_API_KEY, model=ARK_MODEL, timeout=1800):
        if not api_key:
            raise ValueError("请设置ARK_API_KEY环境变量或直接提供API密钥")
            
        self.client = Ark(
            api_key=api_key,
            timeout=timeout  # 深度思考模型耗时较长，设置为30分钟
        )
        self.model = model
    
    def process_document(self, content):
        """调用方舟API处理文档内容，生成标题、摘要和分类"""
        # 构建提示词
        prompt = f"""请仔细分析以下运维文档内容，按照要求进行处理：

1. 生成一个准确反映文档核心内容的标题
2. 撰写一段约100字的简洁摘要，概括文档主要内容
3. 对原文内容进行整理：
   - 适当切分长段落，提升可读性
   - 去除重复信息，保留关键技术点
   - 保持专业术语的准确性
4. 按运维领域标准分类（如：网络设备、服务器、数据库、安全、云平台等）

请严格按照以下格式返回结果（仅返回指定格式内容，不要添加额外说明）：
标题：[生成的标题]
摘要：[生成的摘要]
内容：[处理后的内容]
分类：[确定的分类]

文档内容：
{content[:3000]}
"""
        
        try:
            # 调用方舟SDK的chat completions接口
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是专业的运维文档处理专家，擅长对技术文档进行标准化处理、提炼核心信息和精准分类。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.6,
                max_tokens=2048,
                thinking={
                    "type": "enabled"  # 启用深度思考能力
                }
            )
            
            # 解析响应结果
            if response.choices and len(response.choices) > 0:
                return response.choices[0].message.content
            else:
                print("API返回格式不正确，未找到有效内容")
                return None
                
        except Exception as e:
            print(f"调用方舟API时发生错误: {str(e)}")
            # 打印详细错误信息
            try:
                print(f"错误详情: {e.response.text if hasattr(e, 'response') else str(e)}")
            except:
                pass
            return None

class WordProcessor:
    """Word文档处理工具类"""
    
    @staticmethod
    def read_word_file(file_path):
        """读取Word文档内容"""
        try:
            if file_path.lower().endswith('.docx'):
                doc = docx.Document(file_path)
                full_text = []
                for para in doc.paragraphs:
                    full_text.append(para.text)
                return '\n'.join(full_text)
            elif file_path.lower().endswith('.doc'):
                print(f".doc格式文件 {file_path} 暂不支持，请转换为.docx格式")
                return None
        except Exception as e:
            print(f"读取Word文件 {file_path} 时出错: {str(e)}")
            return None
    
    @staticmethod
    def parse_api_response(response):
        """解析API返回的响应内容"""
        result = {
            "title": "",
            "summary": "",
            "content": "",
            "category": ""
        }
        
        title_match = re.search(r'标题：(.*?)\n', response)
        if title_match:
            result["title"] = title_match.group(1).strip()
            
        summary_match = re.search(r'摘要：(.*?)\n', response, re.DOTALL)
        if summary_match:
            result["summary"] = summary_match.group(1).strip()
            
        content_match = re.search(r'内容：(.*?)\n分类：', response, re.DOTALL)
        if content_match:
            result["content"] = content_match.group(1).strip()
            
        category_match = re.search(r'分类：(.*?)$', response, re.DOTALL)
        if category_match:
            result["category"] = category_match.group(1).strip()
            
        if not result["title"]:
            result["title"] = f"未命名文档_{datetime.now().strftime('%Y%m%d%H%M%S')}"
            
        if not result["category"]:
            result["category"] = "未分类"
            
        return result
    
    @staticmethod
    def write_standard_kb(output_path, title, summary, content):
        """生成标准化的知识库Word文档"""
        try:
            doc = docx.Document()
            
            title_para = doc.add_heading(title, level=0)
            title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            time_para = doc.add_paragraph(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            time_para.alignment = WD_ALIGN_PARAGRAPH.RIGHT
            time_para.style.font.size = Pt(10)
            
            doc.add_paragraph("")
            
            doc.add_heading("摘要", level=1)
            doc.add_paragraph(summary)
            
            doc.add_paragraph("")
            
            doc.add_heading("正文内容", level=1)
            
            paragraphs = content.split('\n')
            for para_text in paragraphs:
                if para_text.strip():
                    para = doc.add_paragraph(para_text.strip())
                    para.style.font.size = Pt(12)
            
            doc.save(output_path)
            return True
        except Exception as e:
            print(f"生成Word文档 {output_path} 时出错: {str(e)}")
            return False

def generate_knowledge_base(input_dir, output_dir):
    """生成运维知识库主函数"""
    try:
        # 初始化方舟API客户端
        ark_client = ArkAPIClient()
    except ValueError as e:
        print(f"初始化失败: {e}")
        return
    
    word_processor = WordProcessor()
    document_info = defaultdict(list)
    
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if file.lower().endswith(('.docx', '.doc')) and not file.startswith('~$'):
                file_path = os.path.join(root, file)
                print(f"处理文件: {file_path}")
                
                project_name = os.path.splitext(file)[0]
                print(f"项目名称: {project_name}")
                
                content = word_processor.read_word_file(file_path)
                if not content:
                    continue
                
                if len(content) > 10000:
                    print("文档内容过长，进行截断处理")
                    content = content[:10000] + "\n...（内容已截断）"
                
                print(f"调用{ARK_MODEL}模型处理内容...")
                api_response = ark_client.process_document(content)
                if not api_response:
                    continue
                
                processed_data = word_processor.parse_api_response(api_response)
                print(f"生成标题: {processed_data['title']}")
                print(f"分类: {processed_data['category']}")
                
                document_info[processed_data['category']].append({
                    'title': processed_data['title'],
                    'project_name': project_name,
                    'summary': processed_data['summary']
                })
                
                category_dir = os.path.join(output_dir, processed_data['category'])
                os.makedirs(category_dir, exist_ok=True)
                
                safe_title = re.sub(r'[\\/*?:"<>|]', '_', processed_data['title'])
                output_file = os.path.join(category_dir, f"{safe_title}.docx")
                
                counter = 1
                while os.path.exists(output_file):
                    output_file = os.path.join(category_dir, f"{safe_title}_{counter}.docx")
                    counter += 1
                
                print(f"生成文档: {output_file}")
                if word_processor.write_standard_kb(
                    output_file,
                    processed_data['title'],
                    processed_data['summary'],
                    processed_data['content']
                ):
                    print("文档生成成功")
                else:
                    print("文档生成失败")
                
                print("-" * 50)
                time.sleep(3)  # 避免请求过于频繁
    
    generate_kb_index(document_info, output_dir)
    print("知识库生成完成")

def generate_kb_index(document_info, output_dir):
    """生成知识库目录"""
    try:
        doc = docx.Document()
        
        title_para = doc.add_heading("运维知识库目录", level=0)
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        time_para = doc.add_paragraph(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        time_para.alignment = WD_ALIGN_PARAGRAPH.RIGHT
        doc.add_paragraph("")
        
        for category in sorted(document_info.keys()):
            doc.add_heading(category, level=1)
            
            for doc_info in document_info[category]:
                p = doc.add_paragraph()
                p.add_run(f"• {doc_info['title']}").bold = True
                p.add_run(f" (项目: {doc_info['project_name']})")
                
                summary_para = doc.add_paragraph(f"  摘要: {doc_info['summary'][:100]}...")
                summary_para.style.font.size = Pt(10)
                doc.add_paragraph("")
        
        index_path = os.path.join(output_dir, "知识库目录.docx")
        doc.save(index_path)
        print(f"知识库目录已生成: {index_path}")
        
    except Exception as e:
        print(f"生成知识库目录时出错: {str(e)}")

if __name__ == "__main__":
    INPUT_DIRECTORY = r"D:\2025运维知识库\解压文档"    # 存放原始Word文档的目录
    OUTPUT_DIRECTORY = r"D:\2025运维知识库\处理生成文档"    # 生成的知识库输出目录
    
    # 创建输出目录
    os.makedirs(OUTPUT_DIRECTORY, exist_ok=True)
    
    # 生成知识库
    generate_knowledge_base(INPUT_DIRECTORY, OUTPUT_DIRECTORY)
    